import toml
from .run_and_save import predict_and_save
from .plot_plane import (
    plt, plot_pos_plane, plot_time_dependent, plot_td_elbow_angle, plot_td_angle, plot_td_argument_angle, 
    plot_arrow_arg_angle)
import argparse
import logging
import os
import pandas as pd


def run_workflow(filename):
    with open(filename, encoding='utf-8', errors='ignore') as fp:
        cfg = toml.load(fp)

    cfgdata = cfg['data']
    cfgana = cfg.get('analysis', dict())
    
    # First check if we need train
    if os.path.exists(cfgdata['keypoint_file']):
        logging.info('Prediction is skipped as keypoint file exists')
    else:
        logging.info('Begin running YOLO')
        predict_and_save(cfgdata['source'], cfgdata.get('model', 'yolo11x-pose.pt'), cfgdata['keypoint_file'])

    kpfile = cfgdata['keypoint_file']
    
    # Now load the data
    df = pd.read_csv(cfgdata['keypoint_file'])
    df = df[df['confident'] >= cfgana.get('confident', 0.8)]
    if (cfgplane := cfg.get('plane')) is not None:
        bg_video_path = cfgdata['source']
        if not cfgplane.get('use_background'):
            bg_video_path = None
        plot_pos_plane(df, cfgana['keypoints'], cfgplane.get('cmap', 'winter'), 
                       bg_video_path, cfgplane.get('background_time', 0))
    
    if (cfgtd := cfg.get('time_dependent')) is not None:
        plot_time_dependent(df, cfgana['keypoints'], title=cfgtd.get('figure_title', cfgdata['keypoint_file']))

    if (cfgeb := cfg.get('td_elbow_angle')) is not None:
        plot_td_elbow_angle(df, cfgeb.get('side', 'both'), title=cfgeb.get('figure_title', cfgdata['keypoint_file']))

    if (cfgarg := cfg.get('td_arrow_arg')) is not None:
        plot_arrow_arg_angle(df, cfgarg.get('draw_side', 'right'), title=cfgarg.get('figure_title', kpfile))

    if (cfgangles := cfg.get('td_angles_raw', [])):
        for ca in cfgangles:
            plot_td_angle(df, ca['center_keypoint'], ca['left_keypoint'], ca['right_keypoint'], ca.get('label', ''))
        plt.show()
    
    if (cfgargs := cfg.get('td_arg_angles_raw', [])):
        for ca in cfgargs:
            plot_td_argument_angle(df, ca['start_keypoint'], ca['end_keypoint'], ca.get('label', None),
                                   ca.get('figure_title', cfgdata['keypoint_file']))
        plt.show()


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument('config_file')
    p = parser.parse_args()
    run_workflow(p.config_file)
